from hyde.hyde_core import HyDERetriever


def main():
    """
    HyDE使用示例
    """
    # 初始化HyDE检索器
    hyde_retriever = HyDERetriever(
        model_url="http://localhost:11434/api/generate",
        model_name="deepseek-r1:1.5b"
    )

    # 添加示例知识库文档
    sample_documents = [
        "Python是一种高级编程语言，由Guido van Rossum于1991年首次发布。它具有简洁的语法和强大的功能。",
        "机器学习是人工智能的一个分支，它使用算法来分析数据、学习模式并做出预测或决策。",
        "深度学习是机器学习的一个子集，它使用多层神经网络来模拟人脑处理信息的方式。",
        "RAG（检索增强生成）是一种结合检索和生成的AI技术，能够提高语言模型的准确性和可靠性。",
        "HyDE（假设性文档嵌入）通过生成假设文档来改进传统RAG方法的检索效果。",
        "向量数据库是专门用于存储和检索高维向量数据的数据库系统，在AI应用中非常重要。"
    ]

    hyde_retriever.add_documents(sample_documents)

    # 执行HyDE检索
    test_queries = [
        #"什么是Python编程语言？",
        #"机器学习和深度学习有什么区别？",
        "如何改进RAG系统的检索效果？"
    ]

    for query in test_queries:
        print(f"\n{'=' * 60}")
        print(f"查询: {query}")
        print('=' * 60)

        result = hyde_retriever.hyde_retrieve(query, top_k=2)

        print(f"\n假设文档:")
        print(f"{result['hypothetical_document']}\n")

        print(f"检索到的文档:")
        for i, doc in enumerate(result['retrieved_documents'], 1):
            print(f"{i}. (相似度: {doc['similarity']:.3f}) {doc['content']}")

        print(f"\n最终答案:")
        print(f"{result['final_answer']}")


if __name__ == "__main__":
    main()